import requests
import time
import json
import webbrowser
from dotenv import load_dotenv
import os

load_dotenv()

API_URL = "https://ai.gitee.com/v1/async/videos/generations"
api_key = os.getenv("api_key")
headers = {
	"Authorization": f"Bearer {api_key}" # 使用 Bearer 令牌进行身份验证
}

def query(payload):
	response = requests.post(API_URL, headers=headers, json=payload)
	return response.json()

def poll_task(task_id):
	status_url = f"https://ai.gitee.com/v1/task/{task_id}"
	timeout = 30 * 60
	retry_interval = 10
	attempts = 0
	max_attempts = int(timeout / retry_interval)
	while attempts < max_attempts:
		attempts += 1
		print(f"Checking task status [{attempts}]...", end="")
		response = requests.get(status_url, headers=headers, timeout=10)
		result = response.json()
		if result.get("error"):
			print('error')
			raise ValueError(f"{result['error']}: {result.get('message', 'Unknown error')}")
		status = result.get("status", "unknown")
		print(status)
		if status == "success":
			if "output" in result and "file_url" in result["output"]:
				file_url = result["output"]["file_url"]
				duration = (result.get('completed_at', 0) - result.get('started_at', 0)) / 1000
				print(f"🔗 Download link: {file_url}")
				print(f"⏱️ Task duration: {duration:.2f} seconds")
				# 在浏览器中打开结果 URL
				webbrowser.open(file_url)
			else:
				print("⚠️ No output URL found")
		elif status in ["failed", "cancelled"]:
			print(f"❌ Task {status}")
		else:
			time.sleep(retry_interval)
			continue
		task_file = f"task_{task_id}.json"
		with open(task_file, "w") as f:
			json.dump(result, f, indent=4)
		print(f"Task was saved to file {task_file}")
		return result
	print(f"⏰ Maximum attempts reached ({max_attempts})")
	return {"status": "timeout", "message": "maximum wait time exceeded"}

if __name__ == "__main__":
	print("Creating task...")
	result = query({
		"model": "Wan2.1-T2V-1.3B",
		"prompt": "一只金鱼在水中游动", # 替换为对目标视频的描述
		"num_inference_steps": 50,
		"num_frames": 81
	})
	task_id = result.get("task_id")
	if not task_id:
		raise ValueError("Task ID not found in the response")
	print(f"Task ID: {task_id}")
	task = poll_task(task_id)
	if task.get("status") == "success":
		# 在这里对任务结果进行处理
		print("Task completed successfully!")
